package common

import (
	"context"
	"github.com/gin-gonic/gin"
	"github.com/opentracing/opentracing-go"
	"github.com/opentracing/opentracing-go/ext"
	zipkinot "github.com/openzipkin-contrib/zipkin-go-opentracing"
)

const (
	TraceID  = "ld-trace-id"
	TraceCtx = "traceCtx"
)

// Trace is gin trace middleware
func Trace() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		var span opentracing.Span
		// get trace from request header
		carrier := opentracing.HTTPHeadersCarrier(ctx.Request.Header)
		trace, err := opentracing.GlobalTracer().Extract(opentracing.HTTPHeaders, carrier)
		if err != nil {
			span = opentracing.StartSpan(ctx.Request.URL.Path)
		} else {
			span = opentracing.StartSpan(ctx.Request.URL.Path, ext.RPCServerOption(trace))
		}
		ext.SpanKind.Set(span, "server")
		ext.Component.Set(span, "net/http")
		ext.HTTPMethod.Set(span, ctx.Request.Method)
		ext.HTTPUrl.Set(span, ctx.Request.URL.String())
		// get traceID
		id := span.Context().(zipkinot.SpanContext).TraceID.String()
		ctx.Writer.Header().Set(TraceID, id)

		// save context
		ctx2 := opentracing.ContextWithSpan(context.Background(), span)
		ctx.Set(TraceCtx, ctx2)

		ctx.Next()
		span.Finish()
	}
}
